# 0导入模块

import tensorflow as tf
import tensorflow.contrib.slim as slim
from tensorflow.contrib.slim import nets
import cv2
import random
import wgs_code.utils as utils

BATCH_SIZE = 8


def get_data(is_train):
    """
    传入是否是训练，返回对应的一批数据
    :param is_train: bool
    :return: image_batch ,shape=[BATCH_SIZE,224,224,3] ,label_batch = [BATCH_SIZE]
    """
    if is_train:
        dir = '../data/train/origin/'
        info_file = '../data/train/all_info.csv'
    else:
        dir = '../data/test/'
        info_file = '../data/test/all_info.csv'

    # 将所有图片的信息解析到all_info这个字典里面
    # all_info = {image_name1: [...], image_name2: [...], image_name3: [...], ...}
    # 其中[...] 为[瑕疵个数, Xmin, Ymin, Xmax, Ymax, 瑕疵名，Xmin, Ymin, Xmax, Ymax, 瑕疵名...]
    all_info = utils.parse_csv_to_dic(info_file)
    # 获取所有image_name
    img_list = list(all_info.keys())

    # 获取一批图片
    image_batch = []
    label_batch = []
    for i in range(BATCH_SIZE):
        # 随机挑选一张图片
        r = random.randint(0, len(img_list) - 1)
        image = cv2.imread(dir + img_list[r])
        image = cv2.resize(image, (224, 224))
        # 获取对应标签并转化为0-1
        label = all_info[img_list[r]][5]
        label = utils.label_type_to_id(label, bin_class=True)
        # 装起来
        image_batch.append(image)
        label_batch.append(label)
    return image_batch, label_batch


# 1定义神经网络的输入,参数和输出,定义前向传播过程
x = tf.placeholder(tf.float32, shape=(None, 224, 224, 3))
y_ = tf.placeholder(tf.int64, shape=(None))

# 使用resnet50模型
net, endpoints = slim.nets.resnet_v2.resnet_v2_50(x)
# 去掉形状为 1 的第 1，2 个索引维度
net = tf.squeeze(net, axis=[1, 2])
# 全连接
logits = slim.fully_connected(net, num_outputs=2)

# 2定义损失函数及反向传播方法
loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y_, logits=logits))
acc = tf.reduce_mean(tf.cast(tf.equal(tf.argmax(logits, 1), y_), dtype=tf.float32))
train_step = tf.train.GradientDescentOptimizer(0.001).minimize(loss)

# 初始化
init_op = tf.global_variables_initializer()

# 03 生成会话,训练STEPS轮
with tf.Session() as sess:
    sess.run(init_op)
    # 训练模型
    STEPS = 3000
    for i in range(STEPS):
        image_batch, label_batch = get_data(True)
        _loss, _acc, _ = sess.run([loss, acc, train_step], feed_dict={x: image_batch, y_: label_batch})
        print("step", i, "loss", _loss, "acc", _acc)
